import traffic_generator
from itertools import product, combinations
import numpy as np

from numpy.random import seed
from numpy.random import rand

'''
Divides the entire
'''
class CliqueTrafficGenerator(traffic_generator.TrafficGenerator):
	def __init__(self, nblocks, ncliques):
		## canot have cliques that are more than the 
		assert(ncliques <= nblocks)
		traffic_generator.TrafficGenerator.__init__(self, nblocks)
		self.ncliques = ncliques
		return

	def get_name(self):
		return "clique{}".format(self.ncliques)

	## returns the probability communication matrix such that all entries in the matrix
	## sums to 1
	def generate_probability_traffic_matrices(self, nsnapshots):
		nblocks_per_clique = self.nblocks/self.ncliques
		leftover_blocks = self.nblocks % self.ncliques
		blocks_in_clique = [None] * self.ncliques
		## assign the blocks to the cliques first
		current_block_id = 0
		traffic_matrices = []
		for clique_id in range(self.ncliques):
			current_clique_num_blocks = nblocks_per_clique
			if leftover_blocks > 0:
				leftover_blocks -= 1
				current_clique_num_blocks += 1
			blocks_in_clique[clique_id] = []
			for _ in range(current_clique_num_blocks):
				blocks_in_clique[clique_id].append(current_block_id)
				current_block_id += 1
		seed(1)
		for _ in range(nsnapshots):
			clique_traffic_distribution = rand(self.ncliques)
			sum_values = sum(clique_traffic_distribution)
			normalized_clique_traffic_distribution = [float(x)/sum_values for x in clique_traffic_distribution]
			##
			tm = np.zeros((self.nblocks, self.nblocks))
			for clique_blocks, traffic_in_clique in zip(blocks_in_clique, normalized_clique_traffic_distribution):
				num_blocks_in_current_clique = len(clique_blocks)
				clique_per_off_diagonal_entry_traffic = traffic_in_clique / (num_blocks_in_current_clique * (num_blocks_in_current_clique - 1))
				for src_block in clique_blocks:
					for dst_block in clique_blocks:
						if src_block != dst_block:
							tm[src_block][dst_block] = clique_per_off_diagonal_entry_traffic
			traffic_matrices.append(tm)
		return traffic_matrices